import os
import sys
import re
import json
from datetime import datetime
from typing import Dict, Any, List

from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from config import COORDINATOR_PROMPT, MAX_ITERATIONS
from agents.base_agent import BaseAgent
from utils.document_builder import DocumentBuilder

from utils.code_parser import extract_function_code, extract_class_code

class DocumentationTask:
    def __init__(self, task_id, task_type, name, file_path, class_name=None):
        self.task_id = task_id
        self.task_type = task_type  # "class" 或 "function"
        self.name = name  # 类名或函数名
        self.file_path = file_path
        self.class_name = class_name  # 如果是方法，这里是所属类名
        self.status = "pending"  # pending, in_progress, completed, failed
        self.result = None
        self.start_time = None
        self.end_time = None
        self.iterations = []  # 存储每次迭代的结果
        self.issues_history = {}  # 追踪已识别的问题及解决状态
        self.improvement_history = []  # 存储历次优化改进点

    
    # 完善_summarize_improvements方法
    def _summarize_improvements(self, prev_doc, current_doc):
        """比较两次迭代之间的文档改进，返回关键改进点"""
        
        # 如果是字符串比较，分析文档的变化
        if isinstance(prev_doc, str) and isinstance(current_doc, str):
            # 简单的长度比较
            length_change = len(current_doc) - len(prev_doc)
            
            # 提取关键部分的改进
            improvements = []
            
            # 检查是否添加了更多的参数描述
            if "参数" in current_doc and "参数" in prev_doc:
                prev_params = re.findall(r'参数[:：](.*?)(?=返回值|$)', prev_doc, re.DOTALL)
                curr_params = re.findall(r'参数[:：](.*?)(?=返回值|$)', current_doc, re.DOTALL)
                
                if prev_params and curr_params:
                    prev_param_count = len(re.findall(r'[-*•]\s*`?\w+`?', prev_params[0]))
                    curr_param_count = len(re.findall(r'[-*•]\s*`?\w+`?', curr_params[0]))
                    
                    if curr_param_count > prev_param_count:
                        improvements.append(f"增加了 {curr_param_count - prev_param_count} 个参数描述")
            
            # 检查是否添加了返回值描述
            if "返回值" not in prev_doc and "返回值" in current_doc:
                improvements.append("添加了返回值描述")
            
            # 检查是否添加了更详细的功能描述
            prev_desc = re.findall(r'描述[:：](.*?)(?=参数|返回值|$)', prev_doc, re.DOTALL)
            curr_desc = re.findall(r'描述[:：](.*?)(?=参数|返回值|$)', current_doc, re.DOTALL)
            
            if prev_desc and curr_desc:
                prev_desc_len = len(prev_desc[0].strip())
                curr_desc_len = len(curr_desc[0].strip())
                
                if curr_desc_len > prev_desc_len * 1.2:  # 描述长度增加20%以上
                    improvements.append("功能描述更加详细")
            
            # 如果没有找到具体改进点，提供一个通用描述
            if not improvements:
                if length_change > 0:
                    improvements.append(f"文档内容增加了 {length_change} 个字符")
                elif length_change < 0:
                    improvements.append(f"文档内容精简了 {-length_change} 个字符")
                else:
                    improvements.append("文档结构优化，但总长度未变")
            
            return improvements
        
        # 如果不是字符串比较，返回空列表
        return ["无法比较非文本内容"]
    
    def add_iteration_result(self, iteration, documentation, score, feedback):
        """添加迭代结果和评分"""
        self.iterations.append({
            'iteration': iteration,
            'documentation': documentation,
            'score': score,
            'feedback': feedback,
            'timestamp': datetime.now()
        })
        
        # 分析这轮迭代解决了哪些问题
        self._analyze_issues(feedback)
        
    def _analyze_issues(self, feedback):
        """分析反馈中的问题并更新问题跟踪状态"""
        # 提取反馈中的问题点
        issues = self._extract_issues(feedback)
        
        # 更新问题历史
        for issue in issues:
            issue_key = issue.lower()[:50]  # 创建问题的键
            if issue_key in self.issues_history:
                # 问题已存在，标记为"待处理"
                self.issues_history[issue_key]['status'] = 'pending'
            else:
                # 新问题
                self.issues_history[issue_key] = {
                    'description': issue,
                    'identified_in': len(self.iterations),
                    'status': 'new',
                    'resolution': None
                }
    
    def _extract_issues(self, feedback):
        """从反馈中提取问题点"""
        issues = []
        # 查找常见的问题标记模式
        patterns = [
            r'需要改进[:：](.*?)(?=\n|$)',
            r'问题[:：](.*?)(?=\n|$)',
            r'缺少(.*?)(?=\n|$)',
            r'建议[:：](.*?)(?=\n|$)'
        ]
        
        for pattern in patterns:
            matches = re.finditer(pattern, feedback, re.MULTILINE)
            for match in matches:
                issues.append(match.group(1).strip())
                
        return feedback
    
    def get_pending_issues(self):
        """获取待解决的问题"""
        return [issue['description'] for key, issue in self.issues_history.items() 
                if issue['status'] in ['new', 'pending']]
    
    def get_optimization_history(self):
        """获取优化历史摘要"""
        history = []
        for i, iter_data in enumerate(self.iterations):
            if i == 0:
                continue  # 跳过第一次迭代
            
            prev_score = self.iterations[i-1]['score'] if i > 0 else 0
            score_change = iter_data['score'] - prev_score
            
            history.append({
                'iteration': iter_data['iteration'],
                'score': iter_data['score'],
                'score_change': score_change,
                'key_improvements': self._summarize_improvements(
                    self.iterations[i-1]['documentation'], 
                    iter_data['documentation']
                )
            })
        
        return history
    
                                   
    def mark_start(self):
        """标记任务开始"""
        self.status = "in_progress"
        self.start_time = datetime.now()
        
    def mark_complete(self, result):
        """标记任务完成"""
        self.status = "completed"
        self.result = result
        self.end_time = datetime.now()
        
    def mark_failed(self, error):
        """标记任务失败"""
        self.status = "failed"
        self.result = error
        self.end_time = datetime.now()
        
    def to_dict(self):
        """将任务转换为字典表示"""
        duration = None
        if self.start_time and self.end_time:
            duration = (self.end_time - self.start_time).total_seconds()
            
        return {
            "task_id": self.task_id,
            "type": self.task_type,
            "name": self.name,
            "file_path": self.file_path,
            "class_name": self.class_name,
            "status": self.status,
            "duration": duration,
            "subtasks": [subtask for subtask in self.subtasks]
        }
        

class CoordinatorAgent(BaseAgent):
    def __init__(self, architect_agent, programmer_agent, reviewer_agent, optimizer_agent):
        super().__init__("协调器")
        self.architect_agent = architect_agent
        self.programmer_agent = programmer_agent
        self.reviewer_agent = reviewer_agent
        self.optimizer_agent = optimizer_agent
        self.document_builder = DocumentBuilder()
        self.tasks = []  # 存储所有文档生成任务

        
        # 加载协调器提示词
        with open(COORDINATOR_PROMPT, 'r', encoding='utf-8') as f:
            coordinator_template = f.read()
        self.prompt = PromptTemplate(
            template=coordinator_template,
            input_variables=["file_paths", "review_feedback", "current_documentation"]
        )
    
    def generate_documentation(self, file_paths, max_iterations=None):
        """协调生成完整文档的过程，先对单个任务迭代优化"""
        if max_iterations is None:
            max_iterations = MAX_ITERATIONS
                
        # 最终文档的构建器
        final_document_builder = DocumentBuilder()
        steps_log = []
        
        print("\n" + "="*50)
        print(f"🚀 开始文档生成过程: {'单文件模式' if len(file_paths)==1 else '多文件模式'}")
        print(f"🔄 每个任务将单独进行最多 {max_iterations} 轮迭代优化")
        print("="*50 + "\n")
        
        # 处理每个文件
        for file_path in file_paths:
            print(f"\n📂 正在处理文件: {file_path}")
            
            # 1. 架构师分析文件结构
            #architect_output = self.architect_agent.analyze_file(file_path)
            #final_document_builder.set_overview_to_submodule(architect_output['overview'])
            
            # 2. 提取文件中的所有类和函数
            classes, functions = self._extract_classes_and_functions(file_path)
            print(f"📊 发现 {len(classes)} 个类和 {len(functions)} 个独立函数")
            
            # 3. 创建所有任务
            all_tasks = []
            
            # 3.1 为每个类创建任务
            for class_name in classes:
                all_tasks.append(DocumentationTask(
                    task_id=len(all_tasks)+1,
                    task_type="class",
                    name=class_name,
                    file_path=file_path
                ))
                
                # 获取类的所有方法
                methods = self._get_class_methods(file_path, class_name)
                for method_name in methods:
                    all_tasks.append(DocumentationTask(
                        task_id=len(all_tasks)+1,
                        task_type="function",
                        name=method_name,
                        file_path=file_path,
                        class_name=class_name
                    ))
            
            # 3.2 为每个独立函数创建任务
            for function_name in functions:
                all_tasks.append(DocumentationTask(
                    task_id=len(all_tasks)+1,
                    task_type="function",
                    name=function_name,
                    file_path=file_path
                ))
            
            # 4. 对每个任务进行迭代优化
            total_tasks = len(all_tasks)
            print(f"\n🔍 共有 {total_tasks} 个文档生成任务需要迭代优化")
            
            for idx, task in enumerate(all_tasks):
                print(f"\n🔍 开始任务 {idx+1}/{total_tasks}: {task.task_type} {task.name}")
                
                # 单任务迭代优化
                optimized_result = self._iterate_single_task(task, max_iterations)
                
                # 将优化后的文档添加到最终文档构建器
                if optimized_result:
                    if task.task_type == "class":
                        self._parse_and_add_class_doc(optimized_result, final_document_builder)
                    else:
                        self._parse_and_add_function_doc(optimized_result, final_document_builder)
        
        # 5. 合并所有任务构建最终文档
        final_documentation = final_document_builder.merge_sections(1)
        final_documentation = self.filter_markdown_code_blocks(final_documentation)
        
        # 保存最终文档
        final_doc_file = os.path.join("./output", f'final_documentation.md')
        with open(final_doc_file, 'w', encoding='utf-8') as f:
            f.write(final_documentation)
        print(f"📄 最终文档已保存至: {final_doc_file}")
        
        print("\n" + "="*50)
        print("🎉 文档生成完成!")
        print("="*50 + "\n")
        
        return {
            'documentation': final_documentation,
            'steps_log': steps_log
        }

    
    def _extract_classes_and_functions(self, file_path):
        """从文件中提取所有类和独立函数"""
        classes = []
        functions = []
        
        # 从架构师的调用图中提取类和函数
        for func_name, info in self.architect_agent.call_graph['function_info'].items():
            if info['file'] == file_path:
                # 如果有类，则添加到类列表
                if 'class' in info and info['class'] and info['class'] not in classes:
                    classes.append(info['class'])
                # 如果没有类，则为独立函数
                elif ('class' not in info or not info['class']) and info['name'] not in functions:
                    functions.append(info['name'])
        
        # 确保类列表唯一
        classes = list(set(classes))
        
        # 如果调用图没有提供足够信息，尝试直接解析文件,可参考code_parser，后续考虑是否合并一下。
        if not classes and not functions:
            try:
                # 读取文件内容
                full_path = os.path.join(self.architect_agent.project_path, file_path)
                with open(full_path, 'r', encoding='utf-8', errors='ignore') as f:
                    content = f.read()
                
                # 使用正则表达式查找类声明
                class_pattern = r'(class|struct)\s+(\w+)'
                for match in re.finditer(class_pattern, content):
                    class_name = match.group(2)
                    if class_name not in classes:
                        classes.append(class_name)
                
                # 使用正则表达式查找全局函数声明
                func_pattern = r'(?:^|\n)(?:static\s+|inline\s+|extern\s+)?(?:void|int|bool|char|double|float|\w+\s*\*|\w+)\s+(\w+)\s*$[^)]*$'
                for match in re.finditer(func_pattern, content):
                    func_name = match.group(1)
                    # 过滤掉常见的C/C++关键字
                    keywords = ['if', 'for', 'while', 'switch', 'return']
                    if func_name not in keywords and func_name not in functions:
                        functions.append(func_name)
            except Exception as e:
                print(f"直接解析文件时出错: {str(e)}")
        
        return classes, functions
    
    def _get_class_methods(self, file_path, class_name):
        """获取类的所有方法"""
        methods = []
        
        # 从架构师的调用图中提取类的方法
        for func_name, info in self.architect_agent.call_graph['function_info'].items():
            if info['file'] == file_path and info.get('class') == class_name:
                methods.append(info['name'])
        
        # 如果调用图没有提供足够信息，尝试直接解析文件--可参考code_parser，后续考虑是否合并一下
        if not methods:
            try:
                # 提取类代码
                class_code = extract_class_code(self.architect_agent.project_path, file_path, class_name)
                
                # 使用正则表达式查找方法声明
                method_pattern = r'(?:virtual\s+|static\s+|inline\s+)?(?:void|int|bool|char|double|float|\w+\s*\*|\w+)\s+(\w+)\s*$[^)]*$'
                for match in re.finditer(method_pattern, class_code):
                    method_name = match.group(1)
                    # 过滤掉常见的C/C++关键字
                    keywords = ['if', 'for', 'while', 'switch', 'return']
                    if method_name not in keywords and method_name not in methods:
                        methods.append(method_name)
            except Exception as e:
                print(f"直接解析类的方法时出错: {str(e)}")
        
        return methods
        
    def _parse_and_add_class_doc(self, class_doc_result, doc_builder=None):
        """解析并添加类文档到文档构建器"""
        if not class_doc_result or 'documentation' not in class_doc_result:
            print("⚠️ 警告: 类文档为空或格式不正确")
            return
        
        # 使用提供的构建器或默认构建器
        builder = doc_builder or self.document_builder
        
        doc_text = class_doc_result['documentation']
        class_name = class_doc_result['name']
        
        # 解析类描述
        description_pattern = r'(?:描述[:：]|类的主要功能[:：])(.*?)(?=(?:字段)|(?:属性)|(?:方法)|(?:\n##)|$)'
        description_match = re.search(description_pattern, doc_text, re.DOTALL)
        description = description_match.group(1).strip() if description_match else doc_text
        
        # 解析属性/字段
        properties = []
        properties_pattern = r'(?:字段|属性)[:：](.*?)(?=(?:方法)|(?:\n##)|$)'
        properties_match = re.search(properties_pattern, doc_text, re.DOTALL)
        
        if properties_match:
            properties_text = properties_match.group(1).strip()
            property_lines = re.split(r'\n+', properties_text)
            
            for line in property_lines:
                if line.strip():
                    property_match = re.search(r'[-*•]?\s*`?(\w+)`?\s*[:：]?\s*(.+)', line.strip())
                    if property_match:
                        prop_name = property_match.group(1)
                        prop_desc = property_match.group(2)
                        properties.append({
                            "name": prop_name,
                            "description": prop_desc,
                            "data_type": ""  # 可以添加类型解析逻辑
                        })
        
        # 添加到指定的文档构建器
        builder.add_class_to_submodule(class_name, description, [], properties)

    def _parse_and_add_function_doc(self, function_doc_result, doc_builder=None):
        """解析并添加函数/方法文档到文档构建器"""
        if not function_doc_result or 'documentation' not in function_doc_result:
            print("⚠️ 警告: 函数文档为空或格式不正确")
            return
        
        # 使用提供的构建器或默认构建器
        builder = doc_builder or self.document_builder
        
        doc_text = function_doc_result['documentation']
        function_name = function_doc_result['name']
        class_name = function_doc_result.get('class_name')
        
        # 解析函数描述
        description_pattern = r'(?:描述[:：]|功能[:：])(.*?)(?=(?:参数)|(?:返回值)|(?:\n##)|$)'
        description_match = re.search(description_pattern, doc_text, re.DOTALL)
        description = description_match.group(1).strip() if description_match else doc_text
        
        # 解析参数
        parameters = []
        params_pattern = r'(?:参数[:：])(.*?)(?=(?:返回值)|(?:\n##)|$)'
        params_match = re.search(params_pattern, doc_text, re.DOTALL)
        
        if params_match:
            params_text = params_match.group(1).strip()
            param_lines = re.split(r'\n+', params_text)
            
            for line in param_lines:
                if line.strip():
                    param_match = re.search(r'[-*•]?\s*`?(\w+)`?\s*(?:$([^)]*)$)?\s*[:：]?\s*(.+)', line.strip())
                    if param_match:
                        param_name = param_match.group(1)
                        param_type = param_match.group(2) if param_match.group(2) else ""
                        param_desc = param_match.group(3)
                        parameters.append({
                            "name": param_name,
                            "description": param_desc,
                            "data_type": param_type
                        })
        
        # 解析返回值
        return_value = {}
        return_pattern = r'(?:返回值[:：])(.*?)(?=\n##|$)'
        return_match = re.search(return_pattern, doc_text, re.DOTALL)
        
        if return_match:
            return_text = return_match.group(1).strip()
            return_type_match = re.search(r'$([^)]*)$', return_text)
            return_value = {
                "description": return_text,
                "data_type": return_type_match.group(1) if return_type_match else ""
            }

        # 如果是类的方法，添加到对应类中，使用指定的文档构建器
        if class_name:
            builder.add_method_to_class(class_name, function_name, description, parameters, return_value)
        else:
            # 独立函数添加到函数列表，使用指定的文档构建器
            builder.add_function_to_submodule(function_name, description, parameters, return_value)

    
    
    def filter_markdown_code_blocks(self, text: str) -> str:

        text = re.sub(r'```markdown\b', '', text)
        def process_text(text):
            matches = list(re.finditer(r'```', text))  # 找到所有的 ```
            if len(matches) > 0:           
                modified = False
                for i in range(len(matches) - 1, -1, -1): 
                    current_pos = matches[i].start()                   
                    # 判断当前反引号后面是否有字母字符
                    next_chars = text[current_pos + 3:current_pos + 4]  
                    if any(char.isalpha() for char in next_chars):  
                        continue  # 直接跳过
                    
                    # 如果当前反引号后面没有字符或为空白字符
                    else :
                        # 查找前一个反引号
                        if i > 0:  # 确保有前一个反引号
                            prev_pos = matches[i - 1].start()  # 直接访问前一个反引号                          
                            # 检查前一个反引号后是否有字母字符
                            next_chars = text[prev_pos + 3:prev_pos + 4]  
                            if any(char.isalpha() for char in next_chars):  
                                continue  # 保留当前反引号
                            else:                             
                                text = text[:current_pos] + text[current_pos+3:]
                                modified = True
                        else:

                            text = text[:current_pos] + text[current_pos+3:]
                            modified = True                 
        
      
            return text
        
            # 调用递归处理函数
        text = process_text(text)
        # 清理多余的空行和空白字符
        text = re.sub(r'\n\s*\n', '\n\n', text).strip()

        return text
    
    def _iterate_single_task(self, task, max_iterations,generate_only = False):
        """对单个任务（类或函数）进行迭代优化，利用历史信息"""
        task_type = task.task_type
        task_name = task.name
        file_path = task.file_path
        class_name = task.class_name
        generate_only = True
        
        print(f"\n{'='*30}")
        print(f"开始对 {task_type} '{task_name}' 进行迭代优化")
        print(f"{'='*30}")
        
        # 初始化
        review_feedback = "初始生成，尚无评审反馈。"
        best_documentation = None
        best_score = 0
        
        for iteration in range(1, max_iterations + 1):
            print(f"\n迭代 {iteration}/{max_iterations} - {task_type} '{task_name}'")
            
            # 执行任务，生成文档
            task.mark_start()
            try:
                if iteration == 1:
                    # 首次迭代，生成初始文档
                    if task_type == "class":
                        result = self.programmer_agent.generate_class_doc(
                            file_path, 
                            task_name
                        )
                    else:
                        result = self.programmer_agent.generate_function_doc(
                            file_path, 
                            task_name, 
                            class_name
                        )
                    
                    # 记录初始文档
                    initial_doc = result['documentation']

                    # 如果标志位为1，直接返回生成的文档
                    if generate_only:
                        print(f"📢 标志位置1，仅生成初始文档，不进行评审和优化。")
                        return result
                else:
                    # 获取上一次迭代的文档和问题
                    previous_doc = task.iterations[-1]['documentation']
                    pending_issues = task.get_pending_issues()
                    
                    # 准备针对性的优化任务信息
                    task_info = {
                        "type": task_type,
                        "name": task_name,
                        "class_name": class_name,
                        "pending_issues": pending_issues,
                        "iteration_history": task.get_optimization_history()
                    }
                    
                    # 指导LLM关注尚未解决的问题
                    focused_review = f"""
    前一轮迭代中未解决的问题:
    {"- " + "\n- ".join(pending_issues) if pending_issues else "无明显问题，继续优化整体质量。"}

    请特别关注解决上述问题，同时提升文档整体质量。
    """
                    
                    # 运行优化
                    optimize_result = self.optimizer_agent.optimize_documentation(
                        previous_doc,
                        focused_review,
                        iteration,
                        task_info
                    )
                    
                    # 更新结果
                    result = {
                        'type': task_type,
                        'name': task_name,
                        'class_name': class_name,
                        'file_path': file_path,
                        'documentation': optimize_result['documentation']
                    }
                
                # 评审本轮文档
                print(f"\n🔍 评审员正在评估 {task_type} '{task_name}' 文档质量...")
                review_result = self.reviewer_agent.review_documentation(result['documentation'],{task_name})
                
                # 记录评分和反馈
                current_score = review_result['score']
                review_feedback = review_result['review']
                
                # 保存本轮迭代结果
                task.add_iteration_result(
                    iteration, 
                    result['documentation'], 
                    current_score, 
                    review_feedback
                )
                
                print(f"\n📊 评审结果 (分数: {current_score}/10):")
                print(f"{review_feedback[:300]}...\n") # 只显示部分反馈以避免输出过长
                
                # 检查优化效果 - 是否有实质改进
                if iteration > 1:
                    prev_score = task.iterations[-2]['score']
                    print(f"📈 分数变化: {prev_score} → {current_score} ({current_score - prev_score:+.1f})")
                    
                    # 检测是否停滞（连续两轮分数无提升）
                    if iteration >= 3:
                        if (task.iterations[-1]['score'] <= task.iterations[-2]['score'] and
                            task.iterations[-2]['score'] <= task.iterations[-3]['score']):
                            print(f"⚠️ 检测到优化停滞，连续3轮无明显改进")
                            # 可以在这里添加突破停滞的策略
                
                # 保存最高分文档
                if current_score > best_score:
                    best_score = current_score
                    best_documentation = result
                    print(f"📈 新的最佳文档! 分数: {best_score}/10")
                
                # 如果评分足够高或达到连续稳定状态，提前结束迭代
                if current_score >= 8.5 or (
                    iteration >= 3 and
                    abs(task.iterations[-1]['score'] - task.iterations[-2]['score']) < 0.2 and
                    abs(task.iterations[-2]['score'] - task.iterations[-3]['score']) < 0.2
                ):
                    print(f"🎉 {task_type} '{task_name}' 文档已达到稳定状态，结束迭代。")
                    break
                    
                task.mark_complete(result)
                    
            except Exception as e:
                print(f"❌ 任务 {task.task_id} 执行失败: {str(e)}")
                task.mark_failed(str(e))
                break
        
        print(f"\n✅ {task_type} '{task_name}' 优化完成，最佳分数: {best_score}/10")
        if task.iterations:
            initial_score = task.iterations[0]['score']
            print(f"📊 优化效果: {initial_score} → {best_score} (提升 {best_score - initial_score:.1f} 分)")
        
        return best_documentation



    def _add_flowchart_to_documentation(self, documentation_result, flowchart_result):
        """将流程图添加到函数文档中"""
        if not documentation_result or not flowchart_result:
            return documentation_result
        
        # 复制文档结果以避免修改原始数据
        doc = documentation_result.copy()
        
        # 提取流程图信息
        function_name = flowchart_result['function_name']
        mermaid_code = flowchart_result['mermaid_code']
        description = flowchart_result['description']
        
        # 创建流程图部分
        flowchart_section = f"\n\n### 流程图\n\n{description}\n\n```mermaid\n{mermaid_code}\n```\n"
        
        # 将流程图添加到文档中
        if 'documentation' in doc:
            doc['documentation'] += flowchart_section
        
        # 存储流程图信息
        doc['flowchart'] = flowchart_result
        
        return doc
